package com.xiaozhi.utils;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.xiaozhi.config.SiliconFlowConfig;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.List;
import java.util.Map;

/**
 * 硅基流重排模型工具类
 *
 * @author knowledge
 */
@Component
@Slf4j
public class SiliconFlowRerankUtil {

    @Autowired
    private SiliconFlowConfig siliconFlowConfig;

    private final HttpClient httpClient;
    private final ObjectMapper objectMapper;

    public SiliconFlowRerankUtil() {
        this.httpClient = HttpClient.newBuilder()
                .connectTimeout(Duration.ofSeconds(30)) // 增加连接超时时间到30秒
                .build();
        this.objectMapper = new ObjectMapper();
    }

    /**
     * 调用硅基流重排模型对文档进行重排序
     *
     * @param query 查询文本
     * @param documents 文档列表
     * @param topN 返回前N个相关文档
     * @return 重排序结果
     * @throws IOException
     * @throws InterruptedException
     */
    public List<RerankResponse.Result> rerank(String query, List<String> documents, int topN) throws IOException, InterruptedException {
        return rerankWithRetry(query, documents, topN, 3); // 添加重试机制，最多重试3次
    }

    /**
     * 带重试机制的重排方法
     *
     * @param query 查询文本
     * @param documents 文档列表
     * @param topN 返回前N个相关文档
     * @param maxRetries 最大重试次数
     * @return 重排序结果
     * @throws IOException
     * @throws InterruptedException
     */
    private List<RerankResponse.Result> rerankWithRetry(String query, List<String> documents, int topN, int maxRetries) throws IOException, InterruptedException {
        int retryCount = 0;
        while (retryCount < maxRetries) {
            try {
                return rerankInternal(query, documents, topN);
            } catch (Exception e) {
                retryCount++;
                if (retryCount >= maxRetries) {
                    log.error("调用硅基流重排模型失败，已重试{}次: {}", maxRetries, e.getMessage(), e);
                    throw e;
                }
                
                log.warn("调用硅基流重排模型失败，正在重试第{}次: {}", retryCount, e.getMessage());
                try {
                    Thread.sleep(1000 * retryCount); // 指数退避策略
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                    throw ie;
                }
            }
        }
        throw new RuntimeException("调用硅基流重排模型失败，已达到最大重试次数");
    }

    /**
     * 实际调用硅基流重排模型的方法
     *
     * @param query 查询文本
     * @param documents 文档列表
     * @param topN 返回前N个相关文档
     * @return 重排序结果
     * @throws IOException
     * @throws InterruptedException
     */
    private List<RerankResponse.Result> rerankInternal(String query, List<String> documents, int topN) throws IOException, InterruptedException {
        // 构造请求
        RerankRequest request = new RerankRequest(
                siliconFlowConfig.getRerankerModel(),
                query,
                documents,
                topN,
                true // 返回文档内容
        );

        String requestBody = objectMapper.writeValueAsString(request);

        // 构造HTTP请求
        HttpRequest httpRequest = HttpRequest.newBuilder()
                .uri(URI.create(siliconFlowConfig.getBaseUrl() + "/rerank"))
                .header("Authorization", "Bearer " + siliconFlowConfig.getApiKey())
                .header("Content-Type", "application/json")
                .timeout(Duration.ofSeconds(60)) // 设置请求超时时间为60秒
                .POST(HttpRequest.BodyPublishers.ofString(requestBody))
                .build();

        // 发送请求
        HttpResponse<String> response = httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());

        if (response.statusCode() != 200) {
            log.error("调用硅基流重排模型失败，状态码: {}，响应: {}", response.statusCode(), response.body());
            throw new RuntimeException("调用硅基流重排模型失败，状态码: " + response.statusCode());
        }

        // 解析响应
        RerankResponse rerankResponse = objectMapper.readValue(response.body(), RerankResponse.class);

        return rerankResponse.getResults();
    }

    /**
     * 重排请求实体类
     */
    @Data
    public static class RerankRequest {
        private String model;
        private String query;
        private List<String> documents;
        private int top_n;
        private boolean return_documents;

        public RerankRequest(String model, String query, List<String> documents, int top_n, boolean return_documents) {
            this.model = model;
            this.query = query;
            this.documents = documents;
            this.top_n = top_n;
            this.return_documents = return_documents;
        }
    }

    /**
     * 重排响应实体类
     */
    /**
     * 重排响应实体类
     */
    @Data
    public static class RerankResponse {
        private String id;          // 接口返回的id字段
        private List<Result> results; // 重排结果列表
        private String model;       // 模型名
        private Usage usage;        // 令牌使用情况（对象类型，替换原long）
        // 新增：匹配API返回的meta字段（用Map兼容任意子结构，后续可改为自定义类）
        private Map<String, Object> meta;

        /**
         * 单个重排结果
         */
        @Data
        public static class Result {
            private int index;              // 原文档索引
            private Document document;      // 关键修改：将String改为Document类
            private double relevance_score; // 相关性分数
        }

        /**
         * Document子实体类：映射接口返回的document对象
         */
        @Data
        public static class Document {
            private String text; // 文档内容（核心字段，你实际需要的文本）
            private String id;      // 文档ID（若接口返回则保留，无则删除）
            // 若接口返回其他子字段（如metadata），可在此处补充
        }

        /**
         * Usage子实体类：映射接口返回的usage对象（解决之前可能的usage解析问题）
         */
        @Data
        public static class Usage {
            private int prompt_tokens;     // 提示令牌数
            private int completion_tokens; // 完成令牌数（重排通常为0）
            private int total_tokens;      // 总令牌数
        }
    }
}